import pandas as pd
import re
import argparse
import csv

def process_tagged_data(input_file, output_file):
    """
    处理带情感标签的CSV文件，将情感标签前的内容保留在post_title中，
    情感标签后的内容提取到description字段中
    
    参数:
        input_file (str): 输入CSV文件路径
        output_file (str): 输出CSV文件路径
    """
    print(f"开始处理文件: {input_file}")
    
    try:
        # 尝试直接读取CSV
        df = pd.read_csv(input_file, encoding='utf-8-sig')
        
        # 创建结果列表
        result_data = []
        
        # 处理每一行数据
        for idx, row in df.iterrows():
            post_id = row['post_id']
            full_text = str(row['post_title'])
            
            # 使用正则表达式寻找中英文情感标签
            sentiment_match = re.search(r'[,，][^,，]*?(正面|负面|中性)[^,，]*?[,，]', full_text)
            
            if sentiment_match:
                # 获取匹配到的完整标签部分
                matched_text = sentiment_match.group(0)
                # 提取情感标签
                sentiment = re.search(r'(正面|负面|中性)', matched_text).group(1)
                
                # 分割位置
                start_pos = sentiment_match.start()
                end_pos = sentiment_match.end()
                
                # 提取标题和描述（处理中英文逗号）
                title = full_text[:start_pos].strip()
                description = full_text[end_pos:].strip()
                
                # 如果描述以逗号开头，移除它
                if description and description[0] in ',，':
                    description = description[1:].strip()
            else:
                # 尝试另一种模式：寻找单独的情感标签
                sentiment_match = re.search(r'(正面|负面|中性)', full_text)
                
                if sentiment_match:
                    sentiment = sentiment_match.group(0)
                    parts = re.split(r'[,，]\s*(正面|负面|中性)\s*[,，]', full_text)
                    
                    if len(parts) >= 2:
                        title = parts[0].strip()
                        description = parts[1].strip()
                    else:
                        # 尝试在情感标签处分割
                        pos = sentiment_match.start()
                        title = full_text[:pos].strip()
                        
                        # 找到情感标签后的第一个逗号
                        rest = full_text[pos + len(sentiment):].strip()
                        comma_pos = -1
                        for i, c in enumerate(rest):
                            if c in ',，':
                                comma_pos = i
                                break
                        
                        if comma_pos != -1:
                            description = rest[comma_pos + 1:].strip()
                        else:
                            description = rest.strip()
                else:
                    # 没有找到情感标签
                    title = full_text
                    sentiment = ""
                    description = ""
            
            # 添加到结果列表
            result_data.append([post_id, title, sentiment, description])
        
        # 创建DataFrame并保存
        result_df = pd.DataFrame(result_data, columns=['post_id', 'post_title', 'sentiment', 'description'])
        result_df.to_csv(output_file, index=False, encoding='utf-8-sig')
        
        print(f"处理完成，共处理 {len(result_df)} 条数据")
        print(f"数据已保存到: {output_file}")
        
        # 显示结果示例
        print("\n处理结果示例（前3条）：")
        print(result_df.head(3).to_string())
        
        return True
    
    except Exception as e:
        print(f"处理文件时出错: {str(e)}")
        
        # 尝试使用更简单的方法
        try:
            print("尝试使用替代方法...")
            result_data = []
            
            with open(input_file, 'r', encoding='utf-8-sig') as f:
                content = f.readlines()
            
            # 跳过标题行
            for line in content[1:]:
                parts = line.strip().split(',', 1)  # 只在第一个逗号处分割
                if len(parts) < 2:
                    continue
                
                post_id = parts[0]
                text = parts[1]
                
                # 按情感标签分割
                title = ""
                sentiment = ""
                description = ""
                
                for label in ["正面", "负面", "中性"]:
                    if label in text:
                        parts = text.split(label, 1)
                        title = parts[0].strip().rstrip(',，').strip()
                        sentiment = label
                        description = parts[1].strip().lstrip(',，').strip()
                        break
                
                if not sentiment:  # 未找到情感标签
                    title = text
                
                result_data.append([post_id, title, sentiment, description])
            
            # 创建DataFrame并保存
            result_df = pd.DataFrame(result_data, columns=['post_id', 'post_title', 'sentiment', 'description'])
            result_df.to_csv(output_file, index=False, encoding='utf-8-sig')
            
            print(f"使用替代方法处理完成，共处理 {len(result_df)} 条数据")
            print(f"数据已保存到: {output_file}")
            
            # 显示结果示例
            print("\n处理结果示例（前3条）：")
            print(result_df.head(3).to_string())
            
            return True
        
        except Exception as e2:
            print(f"替代方法也失败: {str(e2)}")
            return False

def main():
    parser = argparse.ArgumentParser(description='处理带情感标签的CSV文件')
    parser.add_argument('input_file', help='输入CSV文件路径')
    parser.add_argument('-o', '--output', default='taged_data_by_llm.csv', 
                        help='输出CSV文件路径（默认为taged_data_by_llm.csv）')
    
    args = parser.parse_args()
    
    process_tagged_data(args.input_file, args.output)

if __name__ == '__main__':
    main() 